import torch


def load_device(print_flag=False) -> str:
    if torch.cuda.is_available():
        if print_flag:
            print(f"GPU 1: {torch.cuda.get_device_name(0)}")
        return "cuda:0"
    else:
        if print_flag:
            print("CPU")
        return "cpu"


def chk_device():
    # 检查是否有CUDA设备可用
    if torch.cuda.is_available():
        # 获取所有可用的CUDA设备数量
        num_gpus = torch.cuda.device_count()

        # 打印可用的CUDA设备数量
        print(f"Number of available GPUs: {num_gpus}")

        # 获取每个可用的CUDA设备的名字
        for i in range(num_gpus):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
    else:
        print("CPU")
